from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool, StructuredTool, tool
from langchain.agents import initialize_agent, Agent, AgentExecutor, create_react_agent
from langchain.agents import AgentType
from langchain.agents import load_tools
import os
from langchain_community.llms import Tongyi

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"

model = Tongyi(temperature=1)

"""这种定义方式，只能接受一个参数，且不能封装成类，要不然会报少参数"""
class SearchInput(BaseModel):
    query: str = Field(description="search query")

"""使用tool装饰器定义工具，本身就是函数，不需要再定义函数"""
@tool("search-tool", args_schema=SearchInput, return_direct=True)
def search(query: str) -> str:
    """Look up things online."""
    return "我是一个搜索的工具"




"""这种定义方式，可以接受多个参数"""
class CalculatorInput(BaseModel):
    a: str = Field(description="第一个数字")
    b: str = Field(description="第二个数字")


class SortList(BaseModel):
    num: str = Field(description="待排序列表")


def dort_fun(num):
    """Multiply two numbers."""
    return sorted(eval(num))


sorter = StructuredTool.from_function(
    func=dort_fun,  # 工具具体逻辑
    name="sort_num",  # 工具名
    description="排序列表中的数字",  # 工具信息
    args_schema=SortList,  # 工具接受参数信息
    return_direct=True,  # 直接作为工具的输出返回给调用者
    handle_tool_error=True,  # 报错了继续执行，不会吧那些报错行抛出，也可以自定义函数处理，handle_tool_error=函数名
)

"""
zero-shot-react-description
此代理使用ReAct框架，仅基于工具的描述来确定要使用的工具。
可以提供任意数量的工具。
此代理需要为每个工具提供描述
工具只能接受一个参数，不支持多个参数
"""
tools = load_tools(["llm-math"], llm=model)
tools = [sorter] + tools
agent = initialize_agent(tools, model, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  verbose=True)
# 分别调用计算、排序工具进行处理
# print(agent.invoke("计算 2 ** 3 + 4，`[10,4,7]`排一下序,请输出结果"))
# print(agent.invoke("计算 2 ** 3 + 4"))
print(agent.invoke("`[10,4,7]`排一下序,请输出结果"))